#  -*- coding: utf-8 -*-

from pymongo import UpdateOne,ASCENDING, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,calc_negative_diff_dates,multi_computer
from util.database import DB_CONN
import time
from datetime import datetime, timedelta

"""
实现后复权均线的因子计算和保存
"""


class HfqMAFactor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='hfq_ma')

    def computer_single(self,code,is_index, autype, begin_date,end_date):
        dm = DataModule()
        start_time = time.time()
        print('计算后复权均线, %s' % code)

        if begin_date is None:
            begin_date = "1990-12-19"

        df_daily = dm.get_k_data(code, index=is_index, autype=autype,begin_date = None,end_date= end_date)
        if df_daily.index.size > 0:
            df_daily.set_index(['date'], 1, inplace=True)
            update_requests = []

            df_trading_daily = df_daily.loc[df_daily.is_trading == True, :]
            df_trading_daily_copy = df_trading_daily.copy()
            try:
                # 计算MA
                df_trading_daily_copy['ma5'] = round(df_trading_daily_copy['close'].rolling(5).mean(), 2)
                df_trading_daily_copy['ma20'] = round(df_trading_daily_copy['close'].rolling(20).mean(), 2)
                df_trading_daily_copy['ma60'] = round(df_trading_daily_copy['close'].rolling(60).mean(), 2)
                df_trading_daily_copy['ma99'] = round(df_trading_daily_copy['close'].rolling(99).mean(), 2)
                df_trading_daily_copy['ma250'] = round(df_trading_daily_copy['close'].rolling(250).mean(), 2)
                df_trading_daily_copy['ma888'] = round(df_trading_daily_copy['close'].rolling(888).mean(), 2)

                df_target_daily = df_trading_daily_copy.loc[begin_date:end_date]
                for date in df_target_daily.index:
                    update_requests.append(
                        UpdateOne(
                            {'code': code, 'date': date,'index': is_index},
                            {'$set': {'code': code,
                                      'date': date,
                                      'index': is_index,
                                      'ma5': df_target_daily.loc[date]['ma5'],
                                      'ma20': df_target_daily.loc[date]['ma20'],
                                      'ma60': df_target_daily.loc[date]['ma60'],
                                      'ma99': df_target_daily.loc[date]['ma99'],
                                      'ma250': df_target_daily.loc[date]['ma250'],
                                      'ma888': df_target_daily.loc[date]['ma888']}},
                            upsert=True))
            except:
                print('填充后复权均线时发生错误，股票代码：%s 是否指数：%s' % (code,is_index), flush=True)

            if len(update_requests) > 0:
                update_result = self.collection.bulk_write(update_requests, ordered=False)

                end_time = time.time()
                print('填充后复权均线，股票：%s，是否指数：%s,插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                      (code, is_index, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                      flush=True)

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值(后复权的均线)，并保存到数据库中
        5/20/60/99/250/888
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        self.collection.create_index([('code', 1), ('date', 1),('index',1)])

        #获取所有股票
        codes = get_all_codes()
        args = (False,"hfq",begin_date,end_date)
        multi_computer(computer_codes,codes,args)

        # 获取所有指数
        codes = get_all_indexes_date()
        args = (True,None,begin_date,end_date)
        multi_computer(computer_codes,codes,args)


def computer_codes(codes,is_index,autype,begin_date,end_date):
    for code in codes:
        HfqMAFactor().computer_single(code,is_index,autype,begin_date,end_date)


if __name__ == '__main__':
    # 执行因子的提取任务
    #hfq =HfqMAFactor()

    HfqMAFactor().compute('1990-12-19', '2018-09-12')
